import numpy as np
from fastmcp import FastMCP

# 创建MCP服务实例
mcp = FastMCP("Matrix Solver MCP Service")

@mcp.tool()
def solve_linear_equations(A: list, b: list) -> dict:
    """
    使用高斯消元法求解线性方程组 Ax = b
    
    参数:
        A: 系数矩阵 (二维列表)
        b: 常数向量 (一维列表)
        
    返回:
        {
            "status": "success"|"no_solution"|"infinite_solutions",
            "solution": [x1, x2, ...] (如果有解)
        }
    """
    try:
        A_np = np.array(A, dtype=float)
        b_np = np.array(b, dtype=float)
        
        # 高斯消元
        n = len(b)
        for k in range(n):
            # 部分主元选择
            max_row = np.argmax(np.abs(A_np[k:, k])) + k
            if A_np[max_row, k] == 0:
                return {"status": "infinite_solutions"}
                
            # 交换行
            if max_row != k:
                A_np[[k, max_row]] = A_np[[max_row, k]]
                b_np[[k, max_row]] = b_np[[max_row, k]]
            
            # 消元
            for i in range(k+1, n):
                factor = A_np[i, k] / A_np[k, k]
                A_np[i, k:] -= factor * A_np[k, k:]
                b_np[i] -= factor * b_np[k]
        
        # 回代
        x = np.zeros(n)
        for k in range(n-1, -1, -1):
            if A_np[k, k] == 0:
                if b_np[k] == 0:
                    return {"status": "infinite_solutions"}
                else:
                    return {"status": "no_solution"}
            x[k] = (b_np[k] - np.dot(A_np[k, k+1:], x[k+1:])) / A_np[k, k]
            
        return {
            "status": "success",
            "solution": x.tolist()
        }
        
    except Exception as e:
        return {"status": "error", "message": str(e)}

if __name__ == "__main__":
    mcp.run(transport="stdio")